import plotly.express as px
import plotly.graph_objects as go
import pandas as pd
import numpy as np

class FinancialCharts:
    """
    财务数据可视化类
    负责生成财务数据的各种图表
    """
    def __init__(self):
        # 设置默认的图表配置
        self.chart_config = {
            "title": "公司财务结构分析",
            "x_axis_label": "科目分类", 
            "y_axis_label": "金额（亿元）",
            "color_scheme": ["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728"],
            "hover_template": "分类: %{x}<br>金额: %{y:.2f}亿元<br>占比: %{customdata:.1%}"
        }
    
    def create_category_bar_chart(self, data, company_name=None, report_date=None, title=None, periods=None):
        """
        创建按分类的柱状图
        
        参数:
            data (pd.DataFrame或list): 分类后的财务数据，可以是单个DataFrame或DataFrame列表
            company_name (str): 公司名称
            report_date (str): 报告日期
            title (str): 图表标题，如果为None则使用默认标题
            periods (list): 多期数据的期间列表，如果data是列表，则此参数必须提供
            
        返回:
            plotly.graph_objects.Figure: 柱状图对象
        """
        if data is None or data.empty:
            return None
        
        # 准备数据
        # 排除'报表日期'列
        if '报表日期' in data.columns:
            plot_data = data.drop(columns=['报表日期'])
        else:
            plot_data = data.copy()
        
        # 计算总额和占比
        total = plot_data.sum().sum()
        percentages = {col: plot_data[col].sum() / total for col in plot_data.columns}
        
        # 准备绘图数据
        chart_data = []
        for category in plot_data.columns:
            chart_data.append({
                'category': category,
                'amount': plot_data[category].iloc[0] if len(plot_data) > 0 else 0,
                'percentage': percentages[category]
            })
        
        chart_df = pd.DataFrame(chart_data)
        
        # 设置图表标题
        if title is None:
            if company_name and report_date:
                title = f"{company_name} - {report_date} 财务结构分析"
            else:
                title = self.chart_config["title"]
        
        # 创建柱状图
        fig = px.bar(
            chart_df, 
            x='category', 
            y='amount',
            title=title,
            labels={
                'amount': self.chart_config["y_axis_label"], 
                'category': self.chart_config["x_axis_label"]
            },
            color='category',
            color_discrete_sequence=self.chart_config["color_scheme"],
            custom_data=['percentage']
        )
        
        # 设置悬停信息
        fig.update_traces(
            hovertemplate=self.chart_config["hover_template"]
        )
        
        # 设置图表布局
        fig.update_layout(
            xaxis_title=self.chart_config["x_axis_label"],
            yaxis_title=self.chart_config["y_axis_label"],
            showlegend=False,
            height=600,
            width=1000
        )
        
        return fig
    
    def create_multi_period_bar_chart(self, data_list, periods, company_name="", title=None):
        """
        创建多期对比柱状图
        
        参数:
            data_list (list): 多期财务数据列表
            periods (list): 对应的期间列表
            company_name (str): 公司名称
            title (str): 图表标题，如果为None则使用默认标题
            
        返回:
            plotly.graph_objects.Figure: 柱状图对象
        """
        if not data_list or len(data_list) == 0:
            return None
        
        # 准备数据
        chart_data = []
        
        # 如果是单个DataFrame，将其转换为列表格式
        if isinstance(data, pd.DataFrame):
            data_list = [data]
            periods = [report_date if report_date else "当前期间"]
        else:
            data_list = data
            periods = periods if periods else [f"期间{i+1}" for i in range(len(data_list))]
        
        for i, data in enumerate(data_list):
            if data is None or data.empty:
                continue
                
            # 排除'报表日期'列
            if '报表日期' in data.columns:
                plot_data = data.drop(columns=['报表日期'])
            else:
                plot_data = data.copy()
            
            period = periods[i] if i < len(periods) else f"期间{i+1}"
            
            for category in plot_data.columns:
                chart_data.append({
                    'category': category,
                    'amount': plot_data[category].iloc[0] if len(plot_data) > 0 else 0,
                    'period': period
                })
        
        chart_df = pd.DataFrame(chart_data)
        
        # 设置图表标题
        if title is None:
            if company_name:
                title = f"{company_name} - 多期财务结构对比"
            else:
                title = "多期财务结构对比"
        
        # 创建分组柱状图
        fig = px.bar(
            chart_df, 
            x='category', 
            y='amount',
            color='period',
            barmode='group',
            title=title,
            labels={
                'amount': self.chart_config["y_axis_label"], 
                'category': self.chart_config["x_axis_label"],
                'period': '报告期'
            }
        )
        
        # 设置图表布局
        fig.update_layout(
            xaxis_title=self.chart_config["x_axis_label"],
            yaxis_title=self.chart_config["y_axis_label"],
            legend_title="报告期",
            height=600,
            width=1000
        )
        
        return fig

# 测试代码
if __name__ == "__main__":
    # 创建测试数据
    test_data = pd.DataFrame({
        '报表日期': ['2023-12-31'],
        '现金类': [100],
        '应收款类': [50],
        '预付款类': [20],
        '存货类': [80],
        '长期投资类': [120],
        '固定资产类': [200],
        '无形资产&商誉类': [30],
        '短期借款类': [60],
        '应付款类': [40],
        '预收项类': [25]
    })
    
    # 创建图表
    charts = FinancialCharts()
    fig = charts.create_category_bar_chart(test_data, "测试公司", "2023年报")
    
    # 显示图表（在支持的环境中）
    # fig.show()
    
    # 保存图表到HTML文件
    fig.write_html("../data/processed/test_chart.html")
    print("测试图表已保存到 ../data/processed/test_chart.html")